Skip to content

Add DMD2 distillation for Qwen-Image (fastgen) - #1326

Merged
jingyu-ml merged 27 commits into
mainfrom
jingyux/fastgen-1
Jun 8, 2026
Merged

Add DMD2 distillation for Qwen-Image (fastgen)#1326
jingyu-ml merged 27 commits into
mainfrom
jingyux/fastgen-1

Conversation

@jingyu-ml

@jingyu-ml jingyu-ml commented Apr 22, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: New example + new modelopt.torch.fastgen library module.

Adds DMD2 (Distribution Matching Distillation) for Qwen-Image — distilling the base model into a few-step (1–4) generator. Includes the framework-agnostic modelopt.torch.fastgen loss library (DMD pipeline, EMA, optional GAN discriminator) and a NeMo AutoModel–based training example with a mock-data smoke config, a real-data config, and inference / export scripts.

Noted: the example script will be migrated to AutoModel repo

Usage

# Mock-data wiring smoke — runs end-to-end with no dataset to prepare
torchrun --nproc-per-node=8 \
    examples/diffusers/fastgen/dmd2_finetune.py \
    --config examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml

See examples/diffusers/fastgen/README.md for real-data training and inference.

Testing

Unit tests under tests/unit/torch/fastgen/; pre-commit / code-quality clean.

Before your PR is "Ready for review"

  • Backward compatible?: ✅ (new, additive module)
  • Followed CONTRIBUTING.md for any copied code / new deps: ✅
  • New tests added?: ✅
  • Updated Changelog?: N/A

Summary by CodeRabbit

  • New Features

    • Adds a FastGen-based distillation framework (DMD2) with student/fake-score training, EMA support, GAN discriminator branch, inference pipeline, and export utilities.
    • Qwen-Image integration with latent packing and feature-capture for plugin-enabled pipelines.
  • Documentation

    • New README, example configs, and runnable example scripts for Qwen-Image distillation and inference.
  • Tests

    • Comprehensive unit tests covering math parity, gradient routing, plugins, hooks, EMA, and recipe setup.

Loading
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cherry-pick-done Added by bot once PR is cherry-picked to the release branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants